package org.apache.lucene.distributed;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectStreamClass;
import java.io.OutputStream;
import java.net.InetSocketAddress;

import org.apache.lucene.distributed.RMI.ProxyKey;
import org.apache.lucene.distributed.RMI.Server.InternalService;
import org.apache.lucene.distributed.RMIClasses.ClassAccept;

public class RMIServerSession {
  public String id;
  ClassLoader classLoader;
  RMI rmi;
  RMI.Server server;
  ClassService classService;
  ClassAccept classAccept;
  
  public RMIServerSession(String id, RMI rmi, ClassService classService, ClassAccept classAccept) {
    this.id = id;
    this.rmi = rmi;
    server = rmi.getServer();
    this.classService = classService;
    this.classAccept = classAccept;
    classLoader = new SessionClassLoader(Thread.currentThread().getContextClassLoader());
  }
  
  public ObjectInputStream getObjectInputStream(InputStream input) throws IOException {
    return new ServerObjectInputStream(input);
  }
  
  //public ObjectOutputStream getObjectOutputStream(OutputStream out) throws IOException {
  //  return new ServerObjectOutputStream(out);
  //}
  
  public static class RMIServerObjectOutputStream extends ObjectOutputStream {
    RMI.Server server;
    RMI rmi;
    
    public RMIServerObjectOutputStream(OutputStream out, RMI.Server server) throws IOException {
      super(out);
      this.server = server;
      rmi = server.getRMI();
      enableReplaceObject(true);
    }
    
    protected Object replaceObject(Object obj) throws IOException {
      String name = server.getServiceName(obj);
      if (name != null) {
        Class remoteInterface = RMIClasses.getRemoteInterface(obj.getClass());
        InternalService is = server.getInternalService(name);
        return new ProxyKey(name, remoteInterface, new InetSocketAddress(rmi.localhost, rmi.port), is.dgc);
      } else {
        return obj;
      }
    }
  }
  
  public class ServerObjectInputStream extends ObjectInputStream {
    public ServerObjectInputStream(InputStream input) throws IOException {
      super(input);
      enableResolveObject(true);
    }
    
    protected Object resolveObject(Object obj) throws IOException {
      if (obj instanceof ProxyKey) {
        // transform to local rmi proxy
        return rmi.getProxy((ProxyKey) obj);
      } else {
        return obj;
      }
    }

    protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
      String name = desc.getName();
      if (!classAccept.isSpecial(name)) {
        try {
          return super.resolveClass(desc);
        } catch (ClassNotFoundException ex) {
          // could not find class locally, try loading it
          throw new ClassNotFoundException(name);
        }
      }
      Class c = classLoader.loadClass(name);
      if (c == null) 
        throw new ClassNotFoundException(name);
      return c;
    }
  }
  
  public class SessionClassLoader extends ClassLoader {
    public SessionClassLoader(ClassLoader parent) {
      super(parent);
    }

    public InputStream getResourceAsStream(String name) {
      InputStream is = super.getResourceAsStream(name);
      if (is == null) {
        try {
          byte[] bytes = classService.getResource(name);
          if (bytes != null)
            is = new ByteArrayInputStream(bytes);
        } catch (Exception ex) {
          throw new RuntimeException(ex);
        }
      }
      return is;
    }

    protected Class<?> findClass(final String name) throws ClassNotFoundException {
      if (classAccept.isSpecial(name)) {
        byte[] classBytes = null;
        try {
          classBytes = classService.loadClass(id, name);
        } catch (Exception exception) {
          throw new ClassNotFoundException("", exception);
        }
        if (classBytes == null) {
          throw new ClassNotFoundException(name);
          // return super.findClass(name);
        }
        return defineClass(name, classBytes, 0, classBytes.length);
      } else {
        return super.findClass(name);
      }
    }

    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
      Class c = findLoadedClass(name);
      if (c == null) {
        if (classAccept.isSpecial(name)) {
          c = findClass(name);
          if (resolve) {
            resolveClass(c);
          }
          return c;
        } else {
          return super.loadClass(name, resolve);
        }
      }
      return c;
    }
  }
}
